;;===============================================================
;; trie
;; Implements a Trie structure and its related functions, using a mutable Node (Java) object.
;; The mutable object is in use to allow cyclic references in the automaton,
;; a necessity of the Aho-Corasick algorithm.
;;===============================================================
(ns strclj.trie
  (:use [objects.node])
  (:import [objects Node])
)

(defn get-children
  "Get all children of a given node."
  [^Node node]
  (get (node-get-dict node) :children)
)

(defn get-child
  "Returns the specific child for a given key."
  [^Node node key]
  (get (get-children node) key)
)

(defn- update-max-length!
  "Used to update the maximum length of the trie, stored at the root node."
  [^Node root ^long length]
  (node-set-dict! 
    root 
    (assoc (node-get-dict root) 
           :max-length 
           (max length (:max-length (node-get-dict root) 0))
    )
  )
)

(defn- flag-finish-node!
  "We have the starting position of each word at the root. We'd also like to know the ending position.
   That's what this function does."
  [^Node node]
  (node-set-dict! node (assoc (node-get-dict node) :finish true))
)

(defn get-finish-flag
  "Returns true if the node marks the finish of a word, false (nil) otherwise."
  [^Node node]
  (get (node-get-dict node) :finish)
)

(defn- node-add!
  "Adds a node according to a key (char), to the given node. Returns it if already exists."
  [^Node node key]
  (let [child (get-child node key)]
    (if child
      child
      (let [child (node-make (array-map))]
        (node-set-dict! node 
          (assoc-in (node-get-dict node) [:children key] child)
        )
        child
      )
    )
  )
)

(defn- trie-add-aux!
  "Auxilliary function used to recursively add characters to a Trie."
  [^Node node word]
  (if (seq word)
    (recur (node-add! node (first word)) (rest word))
    node
  )
)

(defn trie-add!
  "This function receives a Trie root and a word and uses the trie-add-aux function to add the word to the Trie."
  [^Node trie word]
  (when (seq word)
    (flag-finish-node! (trie-add-aux! trie word))
    (let [length (count word)]
      (update-max-length! trie length))
  )
)

(defn- trie-add-bunch!
  "Adds a bunch of words to the Trie, calling trie-add! each time."
  [^Node trie words]
  (doseq [word words]
    (node-set-dict! trie (trie-add! trie word))
  )
)

(defn trie-make
  "Creates a new trie for the given words."
  [words]
  (let [new-trie (node-make)]
    (trie-add-bunch! new-trie words)
    new-trie
  )
)

(defn trie-get
  "This function retrieves the node matching the given word (prefix) or nil."
  [^Node root word]
  (letfn [(next 
            [^Node node key] 
            (when node (get-child node key)))] 
    (reduce next root word)
  )
)